#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""

"""

class Solution(object):
    def countSubTrees(self, n, edges, labels):
        """
        :type n: int
        :type edges: List[List[int]]
        :type labels: str
        :rtype: List[int]
        """
        self.links = {}     # {num : set()}
        self.labels = labels
        self.result = [0] * n

        self.visited = set()

        self.build_link(edges)
        self.rec(0, set(labels[0]))


        return self.result

    def build_link(self, edges):
        for s, e in edges:
            if s not in self.links:
                self.links[s] = set()
            self.links[s].add(e)

            if e not in self.links:
                self.links[e] = set()
            self.links[e].add(s)


    def rec(self, idx, focus):
        # if self.labels[idx]
        rets = {}

        if idx in self.visited:
            return {}
        else:
            self.visited.add(idx)

        if idx in self.links:
            for e in self.links[idx]:
                new_focus = set(focus)
                new_focus.add( self.labels[e] )
                part_rets = self.rec(e, new_focus)

                for label in part_rets:
                    if label in rets:
                        rets[label] += part_rets[label]
                    else:
                        rets[label] = part_rets[label]

        cur_label = self.labels[idx]
        if cur_label in rets:
            rets[cur_label] += 1
        else:
            rets[cur_label] = 1

        self.result[idx] = rets[self.labels[idx]]

        return rets


def test():
    s = Solution()
    n = 6
    edges = [[0,1],[0,2],[1,3],[3,4],[4,5]]
    labels = "cbabaa"
    ret = s.countSubTrees(n, edges, labels)
    print(ret)

    n = 7
    edges = [[0,1],[1,2],[2,3],[3,4],[4,5],[5,6]]
    labels = "aaabaaa"
    ret = s.countSubTrees(n, edges, labels)
    print(ret)

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
